package core

import (
	"errors"
	"fmt"
	"unsafe"

	"github.com/cogentcore/webgpu/wgpu"
)

// Buffer 异步GPU内存抽象
type Buffer interface {
	ReadAsync(callback func([]float32, error))
	Write(data []float32) error
	Release()
	Size() uint64
	Shape() []uint32
	Raw() *wgpu.Buffer
	Usage() wgpu.BufferUsage
}

type gpuBuffer struct {
	device *wgpu.Device
	buffer *wgpu.Buffer
	queue  *wgpu.Queue
	size   uint64
	shape  []uint32
	pool   *MemoryPool
	usage  wgpu.BufferUsage
}

// NewBuffer 创建GPU buffer
func NewBuffer(ctx *Context, data []float32, shape []uint32, usage wgpu.BufferUsage) (Buffer, error) {
	var totalElements uint64 = 1
	for _, dim := range shape {
		totalElements *= uint64(dim)
	}
	size := totalElements * 4 // float32 = 4字节

	if data != nil && uint64(len(data)) != totalElements {
		return nil, errors.New("data size doesn't match shape")
	}

	// fmt.Printf("NewBuffer: 从内存池获取buffer (size=%d, usage=0x%x)\n", size, uint32(usage))
	poolBuf, err := ctx.MemoryPool().Get(size, usage)
	if err != nil {
		return nil, err
	}

	if len(data) > 0 {
		// fmt.Printf("NewBuffer: 写入初始数据，%d个元素\n", len(data))
		ctx.Queue().WriteBuffer(poolBuf, 0, wgpu.ToBytes(data))
		// fmt.Printf("NewBuffer: 初始数据写入完成\n")
	}

	buf := &gpuBuffer{
		device: ctx.Device(),
		buffer: poolBuf,
		queue:  ctx.Queue(),
		size:   size,
		shape:  shape,
		pool:   ctx.MemoryPool(),
		usage:  usage, // 确保存储了usage
	}

	// fmt.Printf("NewBuffer: 创建完成，usage=0x%x\n", uint32(buf.usage))
	return buf, nil
}

func (b *gpuBuffer) Usage() wgpu.BufferUsage {
	return b.usage
}

func (b *gpuBuffer) Write(data []float32) error {
	if data == nil || uint64(len(data))*4 != b.size {
		return errors.New("invalid data size")
	}
	fmt.Printf("Write: 写入%d个元素\n", len(data))
	b.queue.WriteBuffer(b.buffer, 0, wgpu.ToBytes(data))
	return nil
}

// ReadAsync 异步读取
func (b *gpuBuffer) ReadAsync(callback func([]float32, error)) {
	if callback == nil {
		return
	}

	if b.usage&wgpu.BufferUsageCopySrc == 0 {
		go callback(nil, fmt.Errorf("buffer 缺少 CopySrc usage (当前: 0x%x, 需要: 0x%x)",
			uint32(b.usage), uint32(wgpu.BufferUsageCopySrc)))
		return
	}

	// ✅ 使用staging buffer的实际大小
	stagingDesc := &wgpu.BufferDescriptor{
		Size:  b.size, // 精确大小
		Usage: wgpu.BufferUsageCopyDst | wgpu.BufferUsageMapRead,
		Label: "staging-buffer",
	}

	stagingBuffer, err := b.device.CreateBuffer(stagingDesc)
	if err != nil {
		go callback(nil, fmt.Errorf("staging buffer 创建失败(size=%d): %v", b.size, err))
		return
	}

	encoder, err := b.device.CreateCommandEncoder(nil)
	if err != nil {
		stagingBuffer.Release()
		go callback(nil, fmt.Errorf("encoder 创建失败: %v", err))
		return
	}
	encoder.CopyBufferToBuffer(b.buffer, 0, stagingBuffer, 0, b.size)
	cmdBuf, err := encoder.Finish(nil)
	encoder.Release()
	if err != nil {
		stagingBuffer.Release()
		go callback(nil, fmt.Errorf("command buffer 创建失败: %v", err))
		return
	}

	b.queue.Submit(cmdBuf)
	cmdBuf.Release()

	// ✅ 使用实际buffer size映射
	stagingBuffer.MapAsync(wgpu.MapModeRead, 0, b.size, func(status wgpu.BufferMapAsyncStatus) {
		defer stagingBuffer.Release()

		if status != wgpu.BufferMapAsyncStatusSuccess {
			callback(nil, fmt.Errorf("MapAsync 失败，status=%v，buffer size=%d", status, b.size))
			return
		}

		mappedRange := stagingBuffer.GetMappedRange(0, uint(b.size))
		if len(mappedRange) == 0 {
			callback(nil, fmt.Errorf("映射数据为空，期望size=%d", b.size))
			stagingBuffer.Unmap()
			return
		}

		data := (*[1 << 30]float32)(unsafe.Pointer(&mappedRange[0]))[: b.size/4 : b.size/4]
		result := make([]float32, len(data))
		copy(result, data)
		stagingBuffer.Unmap()
		callback(result, nil)
	})
}

func (b *gpuBuffer) Release() {
	if b.pool != nil {
		b.pool.Put(b.buffer, b.size, b.usage) // 传入usage
	} else {
		b.buffer.Release()
	}
}

func (b *gpuBuffer) Size() uint64      { return b.size }
func (b *gpuBuffer) Shape() []uint32   { return b.shape }
func (b *gpuBuffer) Raw() *wgpu.Buffer { return b.buffer }
